import copy
import pickle

import numpy as np

from racetracks import *
from graph_node import Node
import matplotlib.pyplot as plt
from copy import deepcopy
import math

from dynamic_programming import track_the_best_plan, visualize_the_best_plan

seed = np.random.seed(1234)
graph = {}


# Initializing admissible g value 距离终点的预测值
def init_g_value():
    g_map = deepcopy(race_track).astype(float)
    # Final line G=0
    for point in FINISH_LINE:
        g_map[point[0], point[1]] = 0

    # Start line G=0
    for point in START_LINE:
        g_map[point[0], point[1]] = 0

    # Occupied
    g_map = g_map * 99

    [row_total, col_total] = g_map.shape

    # Zone 1
    # max vel = 4 and at least one step to go to the destination
    for col in range(0, 12):
        g_map[32:, 12 - col - 1] = math.ceil(col / 4)

    # Zone 2
    # Diagonal Heuristic
    goal_x_list = [32, 33, 34]
    goal_y = 11
    g = []
    for px in range(0, row_total - 2):  # row_total
        for py in range(0, col_total):  # col_total -3
            if g_map[px][py] == 0:
                for goal_x in goal_x_list:
                    dx = abs(px - goal_x)
                    dy = abs(py - goal_y)
                    g.append(math.ceil((dx + dy - 2 * min(dx, dy)) / 4) + math.ceil(min(dx, dy) / 4))
                g_map[px, py] = min(g) + 0.01 * abs(px - 32)  # to encourage the racing car to move forward
    return g_map


# 在状态转移图中将每一种状态赋予G值
def g_to_graph(graph, g_map):
    for key in graph:
        px = graph[key].px
        py = graph[key].py
        vx = graph[key].vx
        vy = graph[key].vy

        # give penalty to those cause accidents except final line
        if (px in [32, 33, 34]) and (py == 11):
            continue

        neigh_status = []

        # Give those states causing unavoidable collision  (OUTBOUND, OCCUPIED)
        # penalty

        for neighbor_key in graph[key].next_prob_9:
            # neighbor in START LINE
            if neighbor_key in ['00030000', '00040000', '00050000', '00060000']:
                neigh_status.append(1)
            else:
                neigh_status.append(0)

        neigh_status = np.array(neigh_status)
        if neigh_status.all():
            graph[key].g_value = 99
        else:
            graph[key].g_value = g_map[px][py]

        bound_x_0 = 0
        bound_x_1 = 34

        if 0 <= px < 4:
            bound_y_0 = 3
            bound_y_1 = 6
        elif 4 <= px < 8:
            bound_y_0 = 2
            bound_y_1 = 6
        elif 8 <= px < 12:
            bound_y_0 = 1
            bound_y_1 = 6
        elif 12 <= px < 32:
            bound_y_0 = 0
            bound_y_1 = 6
        else:
            bound_y_0 = 0
            bound_y_1 = 11

        if abs(vx) == 4:
            if abs(px - bound_x_0) <= 6 and abs(px - bound_x_1) <= 6:
                graph[key].g_value = 99
            elif abs(vx) == 3:
                if abs(px - bound_x_0) <= 3 or abs(px - bound_x_1) <= 3:
                    graph[key].g_value = 99
            elif abs(vx) == 2:
                if abs(px - bound_x_0) <= 1 or abs(px - bound_x_1) <= 1:
                    graph[key].g_value = 99
            elif abs(vx) == 1:
                if abs(px - bound_x_0) <= 0 or abs(px - bound_x_1) <= 0:
                    graph[key].g_value = 99

        if abs(vy) == 4:
            if abs(py - bound_y_0) <= 6 or abs(py - bound_y_1) <= 6:
                graph[key].g_value = 99
        elif abs(vy) == 3:
            if abs(py - bound_y_0) <= 3 or abs(py - bound_y_1) <= 3:
                graph[key].g_value = 99
        elif abs(vy) == 2:
            if abs(py - bound_y_0) <= 1 or abs(py - bound_y_1) <= 1:
                graph[key].g_value = 99
        elif abs(vy) == 1:
            if abs(py - bound_y_0) <= 0 or abs(py - bound_y_1) <= 0:
                graph[key].g_value = 99

    return graph


def RTDP(greedy_prob=0.8):
    itr_num = 0
    bellman_error = np.inf
    bellman_error_list = []
    while bellman_error > 0.0001 and itr_num <= 1e4:
        itr_num += 1
        bellman_error = 0.0
        # graph is a dict of Node (=state), graph[state.key] = state
        # state.key is an unique str representing a state (px,py,vx,vy)

        # trajectory
        traj = []
        rand_start = START_LINE[np.random.randint(low=0, high=3, size=1)[0]]
        tmp = Node(rand_start[0], rand_start[1], 0, 0)

        state = graph[tmp.key]
        traj.append(state)

        n = 0
        NUM_INT = 100  # avoid looping
        while state.is_goal is False and n <= NUM_INT:
            n += 1
            neighbor_state_key = []
            neighbor_state_g = []
            for child_idx in range(len(ACTION_SPACE)):
                if ACTION_SPACE[child_idx] == [0, 0]:  # forbid stay at the same position
                    continue
                child_key_9 = state.next_prob_9[child_idx]
                child_9 = graph[child_key_9]

                if child_9 in traj:  # forbid re-visit
                    continue
                else:
                    neighbor_state_key.append(child_key_9)
                    neighbor_state_g.append(child_9.g_value)

            # avoid no state to extend! restart!
            if not neighbor_state_key:
                break

            # greedy or random? 防止一直greedy进入死循环
            if np.random.rand() < greedy_prob:
                min_idx = np.argmin(neighbor_state_g)
                next_state_key = neighbor_state_key[min_idx]
            else:
                rand_idx = np.random.randint(low=0, high=len(neighbor_state_key))
                next_state_key = neighbor_state_key[rand_idx]

            state = graph[next_state_key]
            traj.append(state)
        if n >= NUM_INT:
            print("It seems that there is a loop during greedy policy!")

        for [state_0, state_1] in zip(traj[0:-1], traj[1:]):
            expected_cost_uk_star = 0.9 * (1 + state_1.g_value) + 0.1 * (1 + state_0.g_value)
            current_value = expected_cost_uk_star
            bellman_error += np.linalg.norm(state_0.g_value - current_value)
            state_0.g_value = current_value

        bellman_error_list.append(bellman_error)
        print("{}th iteration: {}".format(itr_num, bellman_error))
    plt.figure()
    x_axis = range(len(bellman_error_list))
    plt.plot(x_axis, bellman_error_list)
    plt.show()


def track_the_best_plan(idx=0):
    start_node = Node(START_LINE[idx][0], START_LINE[idx][1], 0, 0)
    start_key = start_node.key
    state = graph[start_key]
    trajectory = [state]
    # for i in range(grid.shape[0]+grid.shape[1]) a safer condition
    while not state.is_goal:
        value_uk = []
        for child_idx in range(len(ACTION_SPACE)):
            child_key_9 = state.next_prob_9[child_idx]
            child_9 = graph[child_key_9]
            value_uk.append(child_9.g_value)
        child_key = state.next_prob_9[np.argmin(value_uk)]
        state = graph[child_key]
        trajectory.append(state)
        print(state.px, state.py)
    return trajectory


def visualize_the_best_plan(plan, grid_para):
    assert isinstance(plan, list)
    plt.figure(figsize=(4.5, 16))
    plt.pcolor(grid_para, edgecolors='k', linewidths=1)
    plan_len = len(plan)
    plan.append(plan[-1])
    for i in range(plan_len):
        plt.arrow(plan[i].py + 0.5, plan[i].px + 0.5,
                  plan[i + 1].py - plan[i].py, plan[i + 1].px - plan[i].px,
                  color='r', head_width=0.3, head_length=0.1)
    plt.show()


if __name__ == '__main__':
    path = './solution/graph_dp.dat'
    track_map = race_track
    graph = pickle.load(open(path, 'rb'))
    # a  = graph['00030000']
    gMap = init_g_value()
    print(gMap)
    graph = g_to_graph(graph, gMap)

    RTDP(greedy_prob=1.0)

    plan = track_the_best_plan()
    visualize_the_best_plan(plan, track_map)
